Skip to content

fix: Add JAX Container #664

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Closed
wants to merge 0 commits into from
Closed

fix: Add JAX Container #664

wants to merge 0 commits into from
+0 −0

Conversation

Josephrp
Copy link

@Josephrp Josephrp commented Aug 4, 2024

Linked Issue

#662

Benefits of Having This Container

  1. Optimized Performance: JAX uses XLA to compile and run NumPy programs on GPUs, which can significantly speed up numerical computations and machine learning tasks. A container specifically optimized for JAX with CUDA ensures that the environment is configured to leverage GPU acceleration fully.

  2. Reproducibility: Containers encapsulate all dependencies, libraries, and configurations needed to run JAX, ensuring that the environment is consistent across different systems. This is crucial for reproducible research and development.

  3. Ease of Use: Users can easily pull and run the container without worrying about the complex setup required for GPU support and JAX configuration. This reduces the barrier to entry for new users and accelerates development workflows.

  4. Isolation and Security: Containers provide an isolated environment, which enhances security by limiting the impact of potential vulnerabilities. It also avoids conflicts with other software on the host system.

@alexanderankin alexanderankin changed the title feat(new): Add JAX Container fix: Add JAX Container Aug 4, 2024
@alexanderankin alexanderankin added the community-feat feature but its a community module so we wont bump tc core for it label Aug 4, 2024
@Josephrp Josephrp mentioned this pull request Aug 5, 2024
Draft
@alexanderankin
Copy link
Member

github actions having a day:

1s
55s
Prepare all required actions
Getting action download info
Download action repository 'actions/setup-python@v5' (SHA:39cd[1](https://github.com/testcontainers/testcontainers-python/actions/runs/10370387369/job/28708130326?pr=664#step:3:1)4951b08e74b54015e9e001cdefcf80e669f)
Warning: Failed to download action 'https://api.github.com/repos/actions/setup-python/tarball/39cd14951b08e74b54015e9e001cdefcf80e669f'. Error: Response status code does not indicate success: 504 (Gateway Timeout).
Warning: Back off [2](https://github.com/testcontainers/testcontainers-python/actions/runs/10370387369/job/28708130326?pr=664#step:3:2)4.396 seconds before retry.
Warning: Failed to download action 'https://api.github.com/repos/actions/setup-python/tarball/[3](https://github.com/testcontainers/testcontainers-python/actions/runs/10370387369/job/28708130326?pr=664#step:3:3)9cd14951b08e74b54015e9e001cdefcf80e669f'. Error: Response status code does not indicate success: 502 (Bad Gateway).
Warning: Back off 15.372 seconds before retry.

@Josephrp
Copy link
Author

Josephrp commented Sep 5, 2024

#Issue add Nvidia

Santosh — Today at 04:26
Looks great and thanks for adding this! One recommendation is to use the nightly (ghcr.io/nvidia/jax:jax) from JAX Toolbox and you can file any issues there if tests are failing: https://github.com/NVIDIA/JAX-Toolbox.

@Josephrp Josephrp closed this Sep 30, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
community-feat feature but its a community module so we wont bump tc core for it
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@Josephrp @alexanderankin